import torch
from typing import Optional, List, Union, Tuple


class GomokuStateProcessor:
    """五子棋状态处理器 - 支持多步历史信息"""

    def __init__(self, board_size: int = 15, history_steps: int = 4):
        self.board_size = board_size
        self.history_steps = history_steps

    def process_state(self, board_batch: Union[List, torch.Tensor],
                      action_history_batch: Optional[List] = None,
                      cur_player_batch: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        将五子棋状态转换为多通道格式
        
        Args:
            board_batch: 当前棋盘状态 [batch_size, board_size, board_size]
            action_history_batch: 历史动作序列 [batch_size, history_steps] (可选)
            cur_player_batch: 当前玩家 [batch_size] (可选)
        
        Returns:
            处理后的状态张量 [batch_size, 2+history_steps, board_size, board_size]
        """
        # 转换输入格式
        board_batch = self._convert_to_tensor(board_batch)
        board_batch = self._ensure_correct_dimensions(board_batch)

        batch_size = board_batch.shape[0]

        # 创建基本通道
        player_channels = self._create_player_channels(board_batch, cur_player_batch, batch_size)

        # 创建历史通道
        history_channels = self._create_history_channels(action_history_batch, batch_size)

        # 组合所有通道
        all_channels = player_channels + history_channels
        return torch.cat(all_channels, dim=1)

    def _convert_to_tensor(self, board_batch: Union[List, torch.Tensor]) -> torch.Tensor:
        """转换为张量格式"""
        if isinstance(board_batch, list):
            return torch.tensor(board_batch, dtype=torch.float32)
        return board_batch

    def _ensure_correct_dimensions(self, board_batch: torch.Tensor) -> torch.Tensor:
        """确保正确的维度"""
        if len(board_batch.shape) == 2:
            return board_batch.unsqueeze(0).unsqueeze(0)
        elif len(board_batch.shape) == 3:
            return board_batch.unsqueeze(1)
        return board_batch

    def _create_player_channels(self, board_batch: torch.Tensor,
                                cur_player_batch: Optional[torch.Tensor],
                                batch_size: int) -> List[torch.Tensor]:
        """创建玩家通道"""
        state0 = (board_batch > 0).float()
        state1 = (board_batch < 0).float()

        # 根据当前玩家调整通道
        if cur_player_batch is not None:
            for i in range(batch_size):
                if cur_player_batch[i] == -1:
                    state0[i], state1[i] = state1[i].clone(), state0[i].clone()

        return [state0, state1]

    def _create_history_channels(self, action_history_batch: Optional[List],
                                 batch_size: int) -> List[torch.Tensor]:
        """创建历史通道"""
        history_channels = []

        for step in range(self.history_steps):
            channel = torch.zeros((batch_size, 1, self.board_size, self.board_size)).float()
            history_channels.append(channel)

        # 填充历史信息
        if action_history_batch is not None:
            self._fill_history_channels(history_channels, action_history_batch, batch_size)

        return history_channels

    def _fill_history_channels(self, history_channels: List[torch.Tensor],
                               action_history_batch: List, batch_size: int) -> None:
        """填充历史通道信息"""
        for i in range(batch_size):
            history_actions = self._normalize_history_actions(action_history_batch[i])

            for step, action in enumerate(history_actions[-self.history_steps:]):
                if action != -1 and step < self.history_steps:
                    x, y = action // self.board_size, action % self.board_size
                    decay_factor = 1.0 - (step * 0.2)
                    history_channels[step][i][0][x][y] = decay_factor

    def _normalize_history_actions(self, actions: Union[List, Tuple, int]) -> List[int]:
        """标准化历史动作格式"""
        if isinstance(actions, (list, tuple)):
            return list(actions)
        else:
            return [actions]
